// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <cstdint>
#include <kllm/utility/types.h>
#include <kllm/core/lora_adapter.h>
#include <kllm/proto/interface.struct.pb.h>
#include <kllm/utility/cpu.h>
#include <llama.h>
#include <string>
#include <vector>

namespace kllm {


    // sampler parameters
    struct InternalSamplerParams {
        uint32_t seed = LLAMA_DEFAULT_SEED; // the seed used to initialize llama_sampler

        int32_t n_prev             = 64;    // number of previous tokens to remember
        int32_t n_probs            = 0;     // if greater than 0, output the probabilities of top n_probs tokens.
        int32_t min_keep           = 0;     // 0 = disabled, otherwise samplers should return at least min_keep tokens
        int32_t top_k              = 40;    // <= 0 to use vocab size
        float   top_p              = 0.95f; // 1.0 = disabled
        float   min_p              = 0.05f; // 0.0 = disabled
        float   xtc_probability    = 0.00f; // 0.0 = disabled
        float   xtc_threshold      = 0.10f; // > 0.5 disables XTC
        float   typ_p              = 1.00f; // typical_p, 1.0 = disabled
        float   temp               = 0.80f; // <= 0.0 to sample greedily, 0.0 to not output probabilities
        float   dynatemp_range     = 0.00f; // 0.0 = disabled
        float   dynatemp_exponent  = 1.00f; // controls how entropy maps to temperature in dynamic temperature sampler
        int32_t penalty_last_n     = 64;    // last n tokens to penalize (0 = disable penalty, -1 = context size)
        float   penalty_repeat     = 1.00f; // 1.0 = disabled
        float   penalty_freq       = 0.00f; // 0.0 = disabled
        float   penalty_present    = 0.00f; // 0.0 = disabled
        float   dry_multiplier     = 0.0f;  // 0.0 = disabled;      DRY repetition penalty for tokens extending repetition:
        float   dry_base           = 1.75f; // 0.0 = disabled;      multiplier * base ^ (length of sequence before token - allowed length)
        int32_t dry_allowed_length = 2;     // tokens extending repetitions beyond this receive penalty
        int32_t dry_penalty_last_n = -1;    // how many tokens to scan for repetitions (0 = disable penalty, -1 = context size)
        int32_t mirostat           = 0;     // 0 = disabled, 1 = mirostat, 2 = mirostat 2.0
        float   mirostat_tau       = 5.00f; // target entropy
        float   mirostat_eta       = 0.10f; // learning rate
        bool    penalize_nl        = false; // consider newlines as a repeatable token
        bool    ignore_eos         = false;
        bool    no_perf            = false; // disable performance metrics

        std::vector<std::string> dry_sequence_breakers = {"\n", ":", "\"", "*"};     // default sequence breakers for DRY


        std::vector<enum KaiSamplerType> samplers = {
                COMMON_SAMPLER_TYPE_DRY,
                COMMON_SAMPLER_TYPE_TOP_K,
                COMMON_SAMPLER_TYPE_TYPICAL_P,
                COMMON_SAMPLER_TYPE_TOP_P,
                COMMON_SAMPLER_TYPE_MIN_P,
                COMMON_SAMPLER_TYPE_XTC,
                COMMON_SAMPLER_TYPE_TEMPERATURE,
        };

        std::string grammar; // optional BNF-like grammar to constrain sampling

        std::vector<llama_logit_bias> logit_bias; // logit biases to apply

        // print the parameters into a string
        [[nodiscard]] std::string print() const;

        void setup_to_proto(SampleParams &result) const;

        void from_proto(const SampleParams &result);
    };


}  // namespace kllm
